from tests.pytests.helpers.codegen_helpers import (
    compile_chunk,
    assert_code_contains,
    compile_source,
)


# 测试基础的类定义
def test_basic_class():
    src = """
class MyClass:
    y: int = 0

    def __init__(self, x: int):
        self.x = x

    def add(self, a: int, b: int) -> int:
        c = a + b + MyClass.y
        return c + MyClass.get_y()

    @staticmethod
    def get_y() -> int:
        return MyClass.y
"""
    hpp, cpp = compile_source(src)
    assert_code_contains(
        hpp,
        """
        class MyClass {
        public:
            MyClass(int x);
            int add(int a, int b);
            static int get_y();
        public:
            static int y = 0;
            int x;
        };
        """,
    )
    assert_code_contains(
        cpp,
        """
        int MyClass::add(int a, int b) {
            int c = a + b + MyClass::y;
            return c + MyClass::get_y();
        }
        """,
    )


# 测试使用 @dataclass 定义的类，使用 field 定义成员属性
def test_dataclass_field():
    src = """
@dataclass
class MyClass:
    x: int = field(default=0)

    def inc(self, a: int) -> int:
        return a + self.x
"""
    hpp, cpp = compile_source(src)
    assert_code_contains(
        hpp,
        """
        class MyClass {
        public:
            int inc(int a);
        public:
            int x = 0;
        };
        """,
    )
    assert_code_contains(
        cpp,
        """
        int MyClass::inc(int a) {
            return a + this->x;
        }
        """,
    )


def test_static_method_and_field():
    src = """
class MyMath:
    factor: int = 2

    @staticmethod
    def mul(a: int) -> int:
        return a * MyMath.factor
"""
    block = compile_chunk(src)
    assert_code_contains(
        block,
        """
        int MyMath::mul(int a) {
            return a * MyMath::factor;
        }
        """,
    )


def test_class_method():
    src = """
class MyMath:
    factor: int = 2

    def __init__(self, x: int):
        self.x = x

    @staticmethod
    def mul(a: int) -> int:
        m: MyMath = MyMath(1)
        return m.add(a * MyMath.factor)

    def add(self, a: int) -> int:
        return self.x + a
"""
    hpp, cpp = compile_source(src)
    assert_code_contains(
        hpp,
        """
        class MyMath {
        public:
            MyMath(int x);
            static int mul(int a);
            int add(int a);
        public:
            static int factor = 2;
            int x;
        };
        """,
    )
    assert_code_contains(
        cpp,
        """
        MyMath::MyMath(int x) {
            this->x = x;
        }
        int MyMath::mul(int a) {
            MyMath m = MyMath(1);
            return m.add(a * MyMath::factor);
        }
        int MyMath::add(int a) {
            return this->x + a;
        }
        """,
    )
